import os, re
from google.protobuf.descriptor_pb2 import (
    DescriptorProto as MessageDP,
    EnumDescriptorProto as EnumDP,
    FieldDescriptorProto as FieldDP,
    FileDescriptorProto as FileDP,
    ServiceDescriptorProto as ServiceDP,
    OneofDescriptorProto as OneofDP,
)
from google.protobuf.descriptor import Descriptor, FileDescriptor, EnumDescriptor
from google.protobuf.descriptor_pool import DescriptorPool
from .keywords import keywords_in_cc

proto_scalar_types_mapping_cpp = {
    FieldDP.Type.TYPE_INT32: "int32_t",
    FieldDP.Type.TYPE_INT64: "int64_t",
    FieldDP.Type.TYPE_UINT32: "uint32_t",
    FieldDP.Type.TYPE_UINT64: "uint64_t",
    FieldDP.Type.TYPE_SINT32: "int32_t",
    FieldDP.Type.TYPE_SINT64: "int64_t",
    FieldDP.Type.TYPE_FIXED32: "uint32_t",
    FieldDP.Type.TYPE_FIXED64: "uint64_t",
    FieldDP.Type.TYPE_SFIXED32: "int32_t",
    FieldDP.Type.TYPE_SFIXED64: "int64_t",
    FieldDP.Type.TYPE_DOUBLE: "double",
    FieldDP.Type.TYPE_FLOAT: "float",
    FieldDP.Type.TYPE_BOOL: "bool",
    FieldDP.Type.TYPE_STRING: "std::string",
    FieldDP.Type.TYPE_BYTES: "std::string",
}

proto_scalar_types_mapping_ts = {
    FieldDP.Type.TYPE_INT32: "number",
    FieldDP.Type.TYPE_INT64: "bigint",
    FieldDP.Type.TYPE_UINT32: "number",
    FieldDP.Type.TYPE_UINT64: "bigint",
    FieldDP.Type.TYPE_SINT32: "number",
    FieldDP.Type.TYPE_SINT64: "bigint",
    FieldDP.Type.TYPE_FIXED32: "number",
    FieldDP.Type.TYPE_FIXED64: "bigint",
    FieldDP.Type.TYPE_SFIXED32: "number",
    FieldDP.Type.TYPE_SFIXED64: "bigint",
    FieldDP.Type.TYPE_DOUBLE: "number",
    FieldDP.Type.TYPE_FLOAT: "number",
    FieldDP.Type.TYPE_BOOL: "boolean",
    FieldDP.Type.TYPE_STRING: "string",
    FieldDP.Type.TYPE_BYTES: "Uint8Array",
}

proto_scalar_types_mapping_ts_with_ns = {
    FieldDP.Type.TYPE_INT32: "number",
    FieldDP.Type.TYPE_INT64: "number | bigint",
    FieldDP.Type.TYPE_UINT32: "number",
    FieldDP.Type.TYPE_UINT64: "number | bigint",
    FieldDP.Type.TYPE_SINT32: "number",
    FieldDP.Type.TYPE_SINT64: "number | bigint",
    FieldDP.Type.TYPE_FIXED32: "number",
    FieldDP.Type.TYPE_FIXED64: "number | bigint",
    FieldDP.Type.TYPE_SFIXED32: "number",
    FieldDP.Type.TYPE_SFIXED64: "number | bigint",
    FieldDP.Type.TYPE_DOUBLE: "number",
    FieldDP.Type.TYPE_FLOAT: "number",
    FieldDP.Type.TYPE_BOOL: "boolean",
    FieldDP.Type.TYPE_STRING: "string",
    FieldDP.Type.TYPE_BYTES: "Uint8Array",
}


def to_json_name(name: str):
    upper_initial_letter = lambda s: s[0].upper() + s[1:]
    parts = [p for p in name.split("_") if p]
    if name.startswith("_"):
        parts[0] = upper_initial_letter(parts[0])
    return parts[0] + "".join([upper_initial_letter(p) for p in parts[1:]])


def is_scalar(field: FieldDP, allow_bytes: bool = False) -> bool:
    return (
        field.type in proto_scalar_types_mapping_cpp
        and field.label != FieldDP.Label.LABEL_REPEATED
        and (allow_bytes or not is_bytes(field))
    )


def is_string(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_STRING and field.label != FieldDP.Label.LABEL_REPEATED


def is_bytes(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_BYTES and field.label != FieldDP.Label.LABEL_REPEATED


def is_message(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_MESSAGE and field.label != FieldDP.Label.LABEL_REPEATED


def is_enum(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_ENUM and field.label != FieldDP.Label.LABEL_REPEATED


def is_repeated_scalar(field: FieldDP, allow_bytes: bool = False) -> bool:
    return (
        field.type in proto_scalar_types_mapping_cpp
        and field.label == FieldDP.Label.LABEL_REPEATED
        and (allow_bytes or not is_bytes(field))
    )


def is_repeated_string(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_STRING and field.label == FieldDP.Label.LABEL_REPEATED


def is_repeated_bytes(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_BYTES and field.label == FieldDP.Label.LABEL_REPEATED


def is_repeated_message(field: FieldDP, pool: DescriptorPool, allow_map: bool = False) -> bool:
    return (
        field.type == FieldDP.Type.TYPE_MESSAGE
        and field.label == FieldDP.Label.LABEL_REPEATED
        and (allow_map or not is_map(field, pool))
    )


def is_repeated_enum(field: FieldDP) -> bool:
    return field.type == FieldDP.Type.TYPE_ENUM and field.label == FieldDP.Label.LABEL_REPEATED


def is_map(field: FieldDP, pool: DescriptorPool) -> bool:
    if not is_repeated_message(field, pool, allow_map=True):
        return False
    type_name = field.type_name[1:] if field.type_name.startswith(".") else field.type_name
    msg: Descriptor = pool.FindMessageTypeByName(type_name)
    return msg.GetOptions().map_entry


def is_required(field: FieldDP) -> bool:
    return field.label == FieldDP.Label.LABEL_REQUIRED


def is_repeated(field: FieldDP) -> bool:
    return field.label == FieldDP.Label.LABEL_REPEATED


def split_package_and_type_name_of_field_type(field: FieldDP, pool: DescriptorPool):
    field_type_name = field.type_name[1:] if field.type_name.startswith(".") else field.type_name
    pkg_name = ""
    if is_message(field) or is_repeated_message(field, pool):
        msg: Descriptor = pool.FindMessageTypeByName(field_type_name)
        file: FileDescriptor = msg.file
        pkg_name = file.package
    elif is_enum(field) or is_repeated_enum(field):
        enum: EnumDescriptor = pool.FindEnumTypeByName(field_type_name)
        file: FileDescriptor = enum.file
        pkg_name = file.package
    else:
        return "", f"/* unsupported type {field_type_name} */"
    return pkg_name, (field_type_name[len(pkg_name) + 1 :] if pkg_name else field_type_name)


def get_raw_type_name_in_cc(field: FieldDP, pool: DescriptorPool) -> str:
    if is_scalar(field, allow_bytes=True):
        return proto_scalar_types_mapping_cpp[field.type]
    pkg_name, type_name = split_package_and_type_name_of_field_type(field, pool)
    return f"{pkg_name.replace('.', '::')}::{type_name.replace('.', '_')}"


def get_cpp_type(field: FieldDP, pkg_name: str, pool: DescriptorPool) -> str:
    if is_scalar(field, allow_bytes=True) or is_repeated_scalar(field, allow_bytes=True):
        return proto_scalar_types_mapping_cpp[field.type]
    if is_message(field) or is_repeated_message(field, pool):
        return get_raw_type_name_in_cc(field, pool)
    if is_enum(field) or is_repeated_enum(field):
        return "int32_t"
    return f"/* unsupported type {field.type_name} */"


def get_FieldDP_of_value_in_map(field: FieldDP, pool: DescriptorPool) -> FieldDP | None:
    if is_map(field, pool):
        field_type_name = field.type_name[1:] if field.type_name.startswith(".") else field.type_name
        msg: Descriptor = pool.FindMessageTypeByName(field_type_name)
        map_dp = MessageDP()
        msg.CopyToProto(map_dp)
        return map_dp.field[1]
    return None


def get_map_value_cpp_type(field: FieldDP, pool: DescriptorPool) -> tuple[FieldDP | None, str]:
    value_field = get_FieldDP_of_value_in_map(field, pool)
    if value_field is None:
        return (None, f"/* unsupported type {field.type_name} */")
    return (value_field, get_raw_type_name_in_cc(value_field, pool))


def get_ts_type(field: FieldDP, pkg_name: str, pool: DescriptorPool, enable_namespace=False) -> str:
    pkg_name, type_name = split_package_and_type_name_of_field_type(field, pool)
    if is_map(field, pool):
        value_field = get_FieldDP_of_value_in_map(field, pool)
        assert value_field
        return f"{{ [key: string]: {get_ts_type(value_field, pkg_name, pool)} }}"
    scalar_type_mapping_ts = (
        proto_scalar_types_mapping_ts_with_ns if enable_namespace else proto_scalar_types_mapping_ts
    )
    if is_scalar(field, allow_bytes=True):
        return scalar_type_mapping_ts[field.type]
    if is_repeated_scalar(field, allow_bytes=True):
        ele_type = scalar_type_mapping_ts[field.type]
        if "|" in ele_type:
            ele_type = f"({ele_type})"
        return f"{ele_type}[]"
    if is_message(field) or is_enum(field):
        type_name_parts = type_name.split(".")
        if is_message(field):
            type_name_parts[-1] = "I" + type_name_parts[-1]
        return (
            ".".join([*pkg_name.split("."), *type_name_parts]) if enable_namespace else type_name.replace(".", "_")
        )
    if is_repeated_message(field, pool) or is_repeated_enum(field):
        type_name_parts = type_name.split(".")
        if is_repeated_message(field, pool):
            type_name_parts[-1] = "I" + type_name_parts[-1]
        type_name = (
            ".".join([*pkg_name.split("."), *type_name_parts]) if enable_namespace else type_name.replace(".", "_")
        )
        return f"{type_name}[]"
    return f"/* unsupported type {field.type_name} */"


def get_ts_oneof_type(fields: list[FieldDP], pkg_name: str, pool: DescriptorPool, enable_namespace=False) -> str:
    candidates = [
        f'{{case: "{field.json_name}", value: {get_ts_type(field, pkg_name, pool, enable_namespace)}}}'
        for field in fields
    ]
    candidates.append("{case: undefined, value: undefined}")
    return " | ".join(candidates)


def split_fields(msg: MessageDP, include_oneof: bool = False) -> tuple[list[FieldDP], list[FieldDP]]:
    required_fields = []
    optional_fields = []
    for field in msg.field:
        if field.label == FieldDP.Label.LABEL_REQUIRED:
            required_fields.append(field)
        elif include_oneof or not field.HasField("oneof_index"):
            optional_fields.append(field)
    return required_fields, optional_fields


def collect_oneof_fields(msg: MessageDP) -> list[tuple[OneofDP, list[FieldDP]]]:
    oneof_fields = []
    for field in msg.field:
        if field.HasField("oneof_index"):
            if field.oneof_index >= len(oneof_fields):
                oneof_fields.append((msg.oneof_decl[field.oneof_index], [field]))
            else:
                oneof_fields[field.oneof_index][1].append(field)
    return oneof_fields


def get_field_name(field: FieldDP) -> str:
    """Get field name in C++."""
    field_name = field.name.lower()
    if field_name in keywords_in_cc:
        field_name += "_"
    return field_name


def to_napi_value(field: FieldDP, pkg_name: str, pool: DescriptorPool):
    """Convert C++ protobuf field to NAPI value."""
    cpp_type = get_cpp_type(field, pkg_name, pool)
    field_name = get_field_name(field)
    if is_map(field, pool):
        value_fd, value_type = get_map_value_cpp_type(field, pool)
        assert value_fd
        if is_scalar(value_fd, allow_bytes=True) or is_enum(value_fd):
            return f"NapiUtils::CreateMapWithValues<{value_type}/*unused*/>(env, impl->{field_name}())"
        return f"NapiUtils::CreateMapWithValues<{value_type}_Stub>(env, impl->{field_name}())"
    if is_scalar(field) or is_enum(field):
        return f"NapiUtils::CreateScalarValue<{cpp_type}>(env, impl->{field_name}())"
    if is_bytes(field):
        return f"NapiUtils::CreateUInt8Array(env, impl->{field_name}().data(), impl->{field_name}().size())"
    if is_message(field):
        return f"NapiUtils::Wrap<{cpp_type}, {cpp_type}_Stub>(env, impl->mutable_{field_name}())"
    if is_repeated_string(field):
        return f"NapiUtils::CreateArrayWithValues(env, impl->{field_name}(), false)"
    if is_repeated_bytes(field):
        return f"NapiUtils::CreateArrayWithValues(env, impl->{field_name}(), true)"
    if is_repeated_scalar(field, allow_bytes=True) or is_repeated_enum(field):
        return f"NapiUtils::CreateArrayWithValues<{cpp_type}>(env, impl->{field_name}())"
    if is_repeated_message(field, pool):
        return f"NapiUtils::CreateArrayWithValues<{cpp_type}, {cpp_type}_Stub>(env, impl->{field_name}())"
    return "NapiUtils::CreateUndefined(env)"


def from_napi_value(field: FieldDP, napi_value: str, pkg_name: str, pool: DescriptorPool, need_check: bool = False):
    """Convert NAPI value to C++ protobuf field."""
    cpp_type = get_cpp_type(field, pkg_name, pool)
    field_name = get_field_name(field)
    # Check if it's a map field
    if is_map(field, pool):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromMapValue(env, *impl->mutable_{field_name}(), {napi_value}); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value}))) {{ {clear_or_set} }}"
        )
    if is_scalar(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ impl->set_{field_name}(NapiUtils::GetScalarValue<{cpp_type}>(env, {napi_value})); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_enum(field):
        enum_type = get_raw_type_name_in_cc(field, pool)
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ impl->set_{field_name}(static_cast<{enum_type}>(NapiUtils::GetScalarValue<{cpp_type}>(env, {napi_value}))); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_bytes(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ impl->set_{field_name}(NapiUtils::GetUInt8ArrayValue(env, {napi_value})); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_message(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ {cpp_type}_Stub::InitByNapiObj(env, {napi_value}, impl->mutable_{field_name}()); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_repeated_string(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromArrayValue(env, *impl->mutable_{field_name}(), {napi_value}, false); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_repeated_bytes(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromArrayValue(env, *impl->mutable_{field_name}(), {napi_value}, true); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_repeated_scalar(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromArrayValue<{cpp_type}>(env, *impl->mutable_{field_name}(), {napi_value}); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_repeated_message(field, pool):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromArrayValue<{cpp_type}, {cpp_type}_Stub>(env, *impl->mutable_{field_name}(), {napi_value}); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    if is_repeated_enum(field):
        clear_or_set = f"if (NapiUtils::IsUndefined(env, {napi_value})) {{ impl->clear_{field_name}(); }} else {{ NapiUtils::SetFromArrayValue<int32_t>(env, *impl->mutable_{field_name}(), {napi_value}); }}"
        return (
            clear_or_set
            if not need_check
            else f"if (!NapiUtils::IsUndefined(env, {napi_value}) && !NapiUtils::IsNull(env, {napi_value})) {{ {clear_or_set} }}"
        )
    return f"/* unsupported field {field_name} with type {field.type_name} */"


def replace_suffix(name: str, old_suffix: str, new_suffix: str):
    if name.endswith(old_suffix):
        return name[: -len(old_suffix)] + new_suffix
    return name


class NamespaceNode:
    """命名空间节点"""

    def __init__(self, name: str, parent: "NamespaceNode | None" = None):
        self.name = name
        self.data: list[str] = []
        self._sub: dict[str, "NamespaceNode"] = {}
        self.parent = parent

    def sub(self, name: str) -> "NamespaceNode":
        if name in self._sub:
            return self._sub[name]

        node = NamespaceNode(name, self)
        self._sub[name] = node
        return node


def merge_dts(scan_path: str) -> str:
    special_dts_file = {
        "all_in_one.d.ts",  # 用于存储所有的导出符号，由本工具生成
        "Index.d.ts",  # native模块默认的接口声明文件
    }
    # 扫描所有 .d.ts 文件
    root = NamespaceNode("")
    for dir, _, filenames in os.walk(scan_path):
        for filename in filenames:
            if filename.endswith(".d.ts") and filename not in special_dts_file:
                _collect_dts(os.path.join(dir, filename), root)
    # 删除所有 .d.ts 文件及空目录
    for dir, _, filenames in os.walk(scan_path, topdown=False):
        for filename in filenames:
            if filename.endswith(".d.ts") and filename not in special_dts_file:
                os.remove(os.path.join(dir, filename))
        if not any(os.listdir(dir)):
            os.rmdir(dir)

    # 清除不在 namespace 中的文本
    root.data = []
    # 导出 namespace 中的所有数据
    return "".join(_expend2dts(root, ""))


def _expend2dts(node: NamespaceNode, indent: str) -> list[str]:
    # 先展开当前空间的数据，为每行文本增加缩径
    result: list[str] = [indent + line for line in node.data]

    # 再展开子节点的数据
    for sub in node._sub.values():
        result.append(indent)
        if not node.parent:
            # 根节点下的命名空间需要导出
            result.append("export ")

        result.append(f"namespace {sub.name} {{\n")
        result.extend(_expend2dts(sub, indent + "  "))
        result.append(f"{indent}}} // namespace {sub.name}\n")

    return result


def _collect_dts(filePath: str, root: NamespaceNode):
    # 遍历文本
    with open(filePath, "r") as file:
        # namespace 提取所需的正则表达式
        ns_pattern = re.compile(r"^namespace ([^ ]+) {")
        node = root
        for no, line in enumerate(file, 1):
            if line.startswith("namespace"):
                # 新建 namespace 节点
                nsv = ns_pattern.search(line).group(1) # pyright: ignore[reportOptionalMemberAccess]
                node = node.sub(nsv)
            elif line.startswith("} // namespace"):
                # namespace 结束, 回退到上一个命名空间中
                node = node.parent
                if not node:
                    # 根节点结束
                    raise RuntimeError(f"{filePath}:{no}: namespace 节点结束时未找到匹配的 namespace 开始")
                continue
            else:
                # 文本内容
                node.data.append(line)
